Greptile Summary

This PR introduces a persistent grouped MXFP8 quantization kernel with a static grid-stride scheduler. The kernel replaces the previous one-CTA-per-block dispatch with a persistent grid (sm_count × STATIC_PERSISTENT_BLOCKS_PER_SM blocks) that iterates over a virtual work grid, enabling better SM utilization for workloads with large numbers of small tensors (including zero-row tensors in the middle of a group). A TunableConfig struct centralises all performance knobs, ShapeRepresentation is promoted to a compile-time template parameter (eliminating runtime branching inside the hot loop), and a USE_FAST_MATH path is added that uses native BF16/FP16 FMA instructions for higher throughput on SM 10.0+ hardware.

Key changes:

• Persistent scheduler: JobDescriptor / BlockDescriptor abstractions decouple logical block coordinates from physical CTA IDs; advance_to_next_job drives the grid-stride loop.
• Empty-tensor safety: is_job_valid returns true for zero-row/zero-col tensor slots and the outer loop skips them, enabling CUDA-graph-safe group tensors with trailing or interleaved empty slots.
• Fast-math path: New mul_cvt_4x PTX overloads (BF16/FP16 scale × BF16/FP16 data → FP8) added to ptx.cuh; FPx4 gains alignas(4*sizeof(T)) required for the b64 shared-memory loads used in these paths.
• API: nvte_group_quantize_v2 exposes QuantizationConfig (including use_fast_math) to callers.
• Refactoring: ShapeRepresentation consolidated into utils.cuh; two new compile-time dispatch macros (TRANSFORMER_ENGINE_SCALING_TYPE_SWITCH, TRANSFORMER_ENGINE_GROUP_TENSOR_SHAPE_REPRESENTATION_SWITCH) reduce kernel branching at the cost of larger binary size.

Issues found:

• The kernel-internal USE_FAST_MATH && !NON_FP32_CAST_ONLY guard silently returns without writing any output, whereas it should fire a device error so that invalid configurations surface immediately.
• is_job_valid accesses offsets_ptr[tensor_id + 1] for all non-SAME_BOTH_DIMS shapes; this relies on offsets_ptr having num_tensors + 1 entries, an invariant that is not documented or asserted anywhere near the call site.
• The test change to compare outputs as a flat (1, elts_num) shape instead of the original (rows, cols) silently skips scale validation for all rows beyond the first 32.
• STATIC_PERSISTENT_BLOCKS_PER_SM = 24 likely exceeds the hardware limit of ~16 concurrent blocks per SM (2048 threads / 128 threads-per-block), and no compile-time guard enforces that derived TunableConfig combinations remain valid.

Confidence Score: 3/5

• The persistent kernel logic and barrier management are sound, but a silent no-op path for USE_FAST_MATH + float input and a reduced test coverage for scale validation deserve attention before merging.
• The core algorithmic changes are carefully constructed—barrier parity is self-consistent across persistent job iterations, empty-tensor handling is correct, and the double-buffered pipeline is properly sequenced. The fast-math PTX paths are architecture-specific and guarded by CUDA_ARCH checks. The main concerns are: (1) a kernel-level guard silently drops all output instead of erroring for unsupported USE_FAST_MATH + IType combinations; (2) the test suite's switch to flat (1×N) comparison silently stops validating scale rows beyond the first; (3) the offsets_ptr[tensor_id+1] assumption in is_job_valid is undocumented; and (4) STATIC_PERSISTENT_BLOCKS_PER_SM=24 over-provisions the SM for typical hardware thread limits.
• transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh (kernel guard and is_job_valid invariant), tests/cpp/operator/test_cast_mxfp8_grouped.cu (scale comparison coverage).

Important Files Changed

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[nvte_group_quantize_v2] --> B[group_quantize_fwd_helper\nwith QuantizationConfig]
    B --> C{Scaling Type}
    C -->|MXFP8_1D_SCALING| D[mxfp8::group_quantize]
    D --> E[Validate fast_math\nconstraints]
    E --> F{is_single_tensor?}
    F -->|No| G[update_tma_descriptors\n1 thread per tensor block\nSkips empty rows/cols tensors]
    F -->|Yes| H[Use static TMA maps]
    G --> I
    H --> I[Compute work_blocks_X/Y]
    I --> J{PERSISTENT mode?}
    J -->|Yes| K[grid = sm_count × 24 × 1]
    J -->|No| L[grid = work_blocks_X × work_blocks_Y]
    K --> M[group_quantize_mxfp8_kernel\nSHAPE_REP compile-time template\nUSE_FAST_MATH compile-time template]
    L --> M
    M --> N[Init barriers\nBUFFS_NUM=2]
    N --> O{job_finished?}
    O -->|No| P[decode_job → JobDescriptor\nis_job_valid check]
    P --> Q{job_has_work?}
    Q -->|No: empty tensor| R[advance_to_next_job\ncontinue]
    R --> O
    Q -->|Yes| S[decode_block → BlockDescriptor\nfence_acquire_tensormap if new tensor]
    S --> T[Prime pipeline:\nprefetch stage 0]
    T --> U[Process STAGES=4 slices\ncolwise + rowwise per stage\nDouble-buffered TMA]
    U --> V{IS_DBIAS?\nis_single_tensor?}
    V -->|Yes| W[Write partial dbias\nto workspace]
    V -->|No| X
    W --> X[advance_to_next_job]
    X --> O
    O -->|Done| Y[atomicMaxFloat amax_ptr\ndestroy_barriers]
    Y --> Z{IS_DBIAS?}
    Z -->|Yes| AA[grouped_reduce_dbias]
    Z -->|No| AB[Done]
    AA --> AB

_{Last reviewed commit: "Merge branch 'main' ..."}

/te-ci